# encoding:utf-8
# -----------------------------------------------------------
# "ML-Stealer: Stealing Prediction Functionality of Machine Learning Models with Mere Black-Box Access"
# @author: Shijie Wang, 2019.
# ------------------------------------------------------------

"""
Preprocess Dataset
"""
import pandas as pd
import numpy as np
import warnings
from sklearn.model_selection import train_test_split


warnings.filterwarnings("ignore")


def datadeal_adult(filename):
    # get the raw data
    adult_raw = pd.read_csv(filename, header=None, sep=',', na_values=' ?')

    # add header
    adult_raw.rename(columns={0: 'age', 1: 'workclass', 2: 'fnlwgt', 3: 'education', 4: 'education_number',
                             5: 'marriage', 6: 'occupation', 7: 'relationship', 8: 'race', 9: 'sex',
                             10: 'capital_gain', 11: 'capital_loss', 12: 'hours_per_work', 13: 'native_country',
                             14: 'income'}, inplace=True)

    # clean data
    adult_cleaned = adult_raw.dropna()

    # digitalize
    adult_digitization = pd.DataFrame()
    target_columns = ['workclass', 'education', 'marriage', 'occupation', 'relationship', 'race', 'sex',
                      'native_country', 'income']
    for column in adult_cleaned.columns:
        if column in target_columns:
            unique_value = list(enumerate(np.unique(adult_cleaned[column])))
            dict_data = {key: value for value, key in unique_value}
            adult_digitization[column] = adult_cleaned[column].map(dict_data)
        else:
            adult_digitization[column] = adult_cleaned[column]

    for column in adult_digitization:
        adult_digitization[column] = adult_digitization[column].astype(int)
    adult_digitization.to_csv("adult_cleaned.csv")
    print(len(adult_cleaned))

    # construct input and output
    X = adult_digitization[
        ['age', 'workclass', 'fnlwgt', 'education', 'education_number', 'marriage', 'occupation', 'relationship',
         'race', 'sex', 'capital_gain', 'capital_loss', 'hours_per_work', 'native_country']]
    Y = adult_digitization[['income']]

    X_train_1, X_train_2, Y_train_1, Y_train_2 = train_test_split(X, Y, test_size=0.5)

    X_train_1.to_csv("X_adult_train_1.csv", index=None)
    Y_train_1.to_csv("Y_adult_train_1.csv", index=None)
    X_train_1.to_csv("X_adult_train_2.csv", index=None)
    Y_train_2.to_csv("Y_adult_train_2.csv", index=None)


def datadeal_bank(filename):
    # get the raw data
    bank_raw = pd.read_csv(filename, sep=';')

    # add header
    bank_raw.rename(columns={0: 'age', 1: 'job', 2: 'marital', 3: 'education', 4: 'default',
                             5: 'balance', 6: 'housing', 7: 'loan', 8: 'contace', 9: 'day',
                             10: 'month', 11: 'duration', 12: 'campaign', 13: 'pdays',
                             14: 'previous', 15: 'poutcome', 16: 'y'}, inplace=True)

    # digitalize
    bank_digitization = pd.DataFrame()
    target_columns = ['job', 'marital', 'education', 'default', 'housing', 'loan', 'contact',
                      'month', 'poutcome', 'y']
    for column in bank_raw.columns:
        if column in target_columns:
            unique_value = list(enumerate(np.unique(bank_raw[column])))
            dict_data = {key: value for value, key in unique_value}
            bank_digitization[column] = bank_raw[column].map(dict_data)
        else:
            bank_digitization[column] = bank_raw[column]

    for column in bank_digitization:
        bank_digitization[column] = bank_digitization[column].astype(int)
    print(len(bank_raw))

    # construct input and output
    X = bank_digitization[
        ['age', 'job', 'marital', 'education', 'default', 'balance', 'housing', 'loan',
         'contact', 'day', 'month', 'duration', 'campaign', 'pdays', 'previous', 'poutcome'
         ]]
    Y = bank_digitization[['y']]

    X_train_1, X_train_2, Y_train_1, Y_train_2 = train_test_split(X, Y, test_size=0.5)

    X_train_1.to_csv("X_bank_train_1.csv", index=None)
    Y_train_1.to_csv("Y_bank_train_1.csv", index=None)
    X_train_1.to_csv("X_bank_train_2.csv", index=None)
    Y_train_2.to_csv("Y_bank_train_2.csv", index=None)


def datadeal_cancer(filename):
    # get the raw data
    cancer_raw = pd.read_csv(filename, header=None, sep=',', na_values='?')

    # add header
    cancer_raw.rename(columns={0: 'code_number', 1: 'thickness', 2: 'cell_size', 3: 'cell_shape', 4: 'marginal_adhesion',
                               5: 'single_cell_size', 6: 'bare_nuclei', 7: 'bland_chromatin', 8: 'normal_nucleoli',
                               9: 'mitoses', 10: 'class'}, inplace=True)

    # clean data
    cancer_raw = cancer_raw.dropna()

    for column in cancer_raw:
        cancer_raw[column] = cancer_raw[column].astype(int)

    print(len(cancer_raw))
    # construct input and output
    X = cancer_raw[
        ['code_number', 'thickness', 'cell_size', 'cell_shape', 'marginal_adhesion', 'single_cell_size',
         'bare_nuclei', 'bland_chromatin', 'normal_nucleoli', 'mitoses'
         ]]
    Y = cancer_raw[['class']]

    X_train_1, X_train_2, Y_train_1, Y_train_2 = train_test_split(X, Y, test_size=0.5)

    X_train_1.to_csv("X_cancer_train_1.csv", index=None)
    Y_train_1.to_csv("Y_cancer_train_1.csv", index=None)
    X_train_2.to_csv("X_cancer_train_2.csv", index=None)
    Y_train_2.to_csv("Y_cancer_train_2.csv", index=None)


if __name__ == '__main__':
    datadeal_adult('./datasets/adult/adult_train.csv')
    datadeal_bank('./datasets/bank/bank-full.csv')
    datadeal_cancer('./datasets/cancer/cancer.csv')
